from tqdm import tqdm
import fasttext
from config import conf
import jieba
from common import get_id2label, ensure_pdir_exist


def do_cut(text, cut_type):
    """
    cut_type: 按字符或按词切分处理，'char' or 'token'
    """
    if cut_type == 'token':
        return ' '.join(jieba.lcut(text))
    if cut_type == 'char':
        return ' '.join(text)
    raise ValueError('type error')


def preprocess(path, target_path, cut_type):
    preprocessed_data = []
    id2label = get_id2label()
    with open(path, 'r', encoding='utf-8') as f:
        for line in tqdm(f.readlines(), desc='开始处理...'):
            text, label = line.strip().split('\t')
            cut = do_cut(text, cut_type)
            new_line = f'__label__{id2label[int(label)]} {cut}'
            preprocessed_data.append(new_line)
    print(f'处理后数据:\n{preprocessed_data[:5]}')
    ensure_pdir_exist(target_path)
    with open(target_path, 'w', encoding='utf-8') as f:
        f.write('\n'.join(preprocessed_data))
    print(f'处理数据完成，保存到{target_path}')


def train_with_fasttext(input_path, model_path, test_path, cut_type,
                        auto_tune=None,
                        val_path=None,
                        autotune_duration=None,
                        ):
    if not auto_tune:
        model = fasttext.train_supervised(input_path)
    else:
        model = fasttext.train_supervised(
            input_path,
            autotuneValidationFile=val_path,
            autotuneDuration=autotune_duration
        )
    ensure_pdir_exist(model_path)
    model.save_model(model_path)
    text = '美财政部计划下周发行660亿美元国债'
    input = do_cut(text, cut_type)
    output = model.predict(input)
    print(f'预测结果 text={text} label={output[0]}')
    print(f'模型词表：{model.get_words()[:10]}')
    subwords = model.get_subwords('你好')
    print(f'查看子词：{subwords}')
    num_samples, acc, prec = model.test(test_path)
    print(f'测试结果 num_samples={num_samples} acc={acc} prec={prec}')


def do_preproc():
    preprocess(conf.dev_path, conf.fasttext_config.dev_char_path, 'char')
    preprocess(conf.test_path, conf.fasttext_config.test_char_path, 'char')
    preprocess(conf.train_path, conf.fasttext_config.train_char_path, 'char')
    preprocess(conf.dev_path, conf.fasttext_config.dev_token_path, 'token')
    preprocess(conf.test_path, conf.fasttext_config.test_token_path, 'token')
    preprocess(conf.train_path, conf.fasttext_config.train_token_path, 'token')


def do_train():
    # train_with_fasttext(
    #     conf.fasttext_config.train_char_path,
    #     conf.fasttext_config.char_model,
    #     conf.fasttext_config.test_char_path,
    #     'char',
    # )
    # # 测试结果 num_samples=10000 acc=0.8763 prec=0.8763
    # train_with_fasttext(
    #     conf.fasttext_config.train_token_path,
    #     conf.fasttext_config.token_model,
    #     conf.fasttext_config.test_token_path,
    #     'token',
    # )
    # # 测试结果 num_samples=10000 acc=0.9071 prec=0.9071
    # train_with_fasttext(
    #     conf.fasttext_config.train_char_path,
    #     conf.fasttext_config.char_auto_model,
    #     conf.fasttext_config.test_char_path,
    #     'char',
    #     auto_tune=True,
    #     val_path=conf.fasttext_config.dev_char_path,
    #     autotune_duration=20 * 60
    # )
    # # 测试结果 num_samples=10000 acc=0.923 prec=0.923
    train_with_fasttext(
        conf.fasttext_config.train_token_path,
        conf.fasttext_config.token_auto_model,
        conf.fasttext_config.test_token_path,
        'token',
        auto_tune=True,
        val_path=conf.fasttext_config.dev_token_path,
        autotune_duration=15 * 60
    )
    # 测试结果 num_samples=10000 acc=0.9218 prec=0.9218
    pass


_model = None


def get_model():
    global _model
    if _model is None:
        _model = fasttext.load_model(conf.fasttext_config.token_auto_model)
    return _model


def predict(text):
    return get_model().predict(do_cut(text, 'token'))[0][0][9:]


def _test_predict():
    text = '港企在港成功首发13.8亿元人民币债券'
    print(f'text={text} result={predict(text)}')
    text = '文化部酝酿推面向未成年人网游企业夜间断网'
    print(f'text={text} result={predict(text)}')


if __name__ == '__main__':
    # do_preproc()
    do_train()
    _test_predict()
    pass
